# -*- coding: utf-8 -*-
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *
import cv2
import json,re,requests
import os
import random
from PyQt5 import QtWidgets
import shutil
import threading
from collections import OrderedDict
from lib.drawBox import DrawBox
from lib.filedock import  fileDock
from lib.scroll import Scroll
from lib.detector import *
from lib.lshape import Shape
from lib.wizard import MagicWizard
import AppResource
from lib.canvas import Canvas

from lib.pconf import Config, sys_cfg

extensions = ['*.%s ' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]

class LBox(QMainWindow):
    CREATE, EDIT = list(range(2))
    def __init__(self, path=None, pathSave=None):
        super(LBox, self).__init__()
        self.resize(1300, 800)

        self.directory = []
        if path is not None:
            self.directory.append(path)

        self.zoom = 100
        self.enableCheckSave = False
        self.isDetecting = False
        self.cfg = Config()
        self.cfg.parse_arguments(os.path.join(sys_cfg["system_rootpath"], "cfg/conf.yaml"))
        self.cfg["system_rootpath"] = sys_cfg["system_rootpath"]

        self.initUI()
        self.createAction()
        self.createToolbar()
        self.createMenu()
        self.setCenter()

    def setCenter(self):
        screen = QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move(int((screen.width() - size.width()) / 2) , int((screen.height() - size.height()) / 2))

    def initUI(self):
        widget = QWidget()
        self.grid = QGridLayout()
        widget.setLayout(self.grid)

        # create main show window
        self.createCenter(self.grid)
        self.label_status = QLabel()
        self.grid.addWidget(self.label_status, 2,0)

        self.pbar = QProgressBar(self)
        self.pbar.hide()
        self.grid.addWidget(self.pbar, 2, 1)

        self.barValue = BarStatusWorker()
        self.barValue.progressBarValue.connect(self.set_par)


        self.grid.setRowStretch(0, 1)
        self.grid.setRowStretch(1, 50)
        self.grid.setRowStretch(2, 1)

        # create image list
        self.filedock1 = fileDock(self, '图片列表:', self.directory)
        self.filedock1.fileListWidget.itemClicked.connect(self.filedock1.fileitemClicked)
        self.filedock1.next.clicked.connect(self.filedock1.openNextImg)
        self.filedock1.prev.clicked.connect(self.filedock1.openPrevImg)

        self.filedock1.newCanvas.connect(self.newCanvas)
        self.filedock1.new_dir.clicked.connect(self.fileDockerAddFile)

        self.setWindowTitle('深度学习检测工具')
        self.setCentralWidget(widget)
        self.addDockWidget(Qt.RightDockWidgetArea, self.filedock1.filedock)

    def set_par(self, i):
        self.pbar.setValue(i)

    def createCenter(self, grid):
        self.v1 = []
        self.image1 = DrawBox(self.v1, self.update, self)
        self.scroll1 = Scroll(self, self.image1)
        self.image1.zoomRequest.connect(self.zoomRequest1)
        grid.addWidget(self.scroll1.scroll, 1, 0, 1, 2)

    def newShape(self):
        items = self.uniqLabelList.selectedItems()
        text = None
        if items:
            text = items[0].data(Qt.UserRole)
        flags = {}
        group_id = None
        if self._config["display_label_popup"] or not text:
            previous_text = self.labelDialog.edit.text()
            text, flags, group_id = self.labelDialog.popUp(text)
            if not text:
                self.labelDialog.edit.setText(previous_text)

        if text and not self.validateLabel(text):
            self.errorMessage(
                self.tr("Invalid label"),
                self.tr("Invalid label '{}' with validation type '{}'").format(
                    text, self._config["validate_label"]
                ),
            )
            text = ""
        if text:
            self.labelList.clearSelection()
            shape = self.canvas.setLastLabel(text, flags)
            shape.group_id = group_id
            self.addLabel(shape)
            self.actions.editMode.setEnabled(True)
            self.actions.undoLastPoint.setEnabled(False)
            self.actions.undo.setEnabled(True)
            self.setDirty()
        else:
            self.canvas.undoLastLine()
            self.canvas.shapesBackups.pop()

    def scrollRequest(self, delta, orientation):
        units = -delta * 0.1  # natural scroll
        bar = self.scrollBars[orientation]
        value = bar.value() + bar.singleStep() * units
        self.setScroll(orientation, value)

    def zoomRequest(self, delta, pos):
        canvas_width_old = self.canvas.width()
        units = 1.1
        if delta < 0:
            units = 0.9
        self.addZoom(units)

        canvas_width_new = self.canvas.width()
        if canvas_width_old != canvas_width_new:
            canvas_scale_factor = canvas_width_new / canvas_width_old

            x_shift = round(pos.x() * canvas_scale_factor) - pos.x()
            y_shift = round(pos.y() * canvas_scale_factor) - pos.y()

            self.setScroll(
                Qt.Horizontal,
                self.scrollBars[Qt.Horizontal].value() + x_shift,
            )
            self.setScroll(
                Qt.Vertical,
                self.scrollBars[Qt.Vertical].value() + y_shift,
            )

    def update(self):
        self.image1.update()

    def zoomRequest1(self, value):
        self.scroll1.zoomRequest(value)
        self.image1.epsilon = self.image1.base_epsilon / (self.image1.scale * 10)

    def createAction(self):
        self.exitAction = QAction('&退出', self)
        self.exitAction.setShortcut('Ctrl+Q')
        self.exitAction.triggered.connect(qApp.quit)

        self.readAction = QAction(QIcon(":/icons/newProject.png"), "& 新 建 \n项 目 ", self)
        self.readAction.triggered.connect(self.newProject)

        self.loadProjectAct = QAction(QIcon(":/icons/loadProject.png"), "& 加 载 \n项 目 ", self)
        self.loadProjectAct.triggered.connect(self.loadProject)

        self.saveAction = QAction(QIcon(":/icons/save.png"), '& 保 存 \n项 目 ', self)
        self.saveAction.setShortcut('Ctrl+S')
        self.saveAction.triggered.connect(self.save)

        self.labelObject = QAction(QIcon(":/icons/label.png"), '& 标 注 \n图 片 ', self)
        self.labelObject.triggered.connect(self.labelImg)

        self.detectObject = QAction(QIcon(":/icons/detect.png"), '& 模 型 \n预 测 ', self)
        self.detectObject.triggered.connect(self.runDetector)

        self.trainDetector = QAction(QIcon(":/icons/train.png"), '& 训 练 \n模 型 ', self)
        self.trainDetector.triggered.connect(self.runTrainer)


    def createToolbar(self):
        toolbar = QToolBar('文件')
        # toolbar.setOrientation(Qt.Vertical)

        self.addToolBar(Qt.LeftToolBarArea, toolbar)
        toolbar.addAction(self.readAction)
        toolbar.addAction(self.loadProjectAct)
        toolbar.addAction(self.saveAction)
        toolbar.addAction(self.labelObject)
        toolbar.addAction(self.detectObject)
        toolbar.addAction(self.trainDetector)
        toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        # toolbar.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)


    def createMenu(self):
        menubar = self.menuBar()
        fileMenu = menubar.addMenu('&文件')
        fileMenu.addAction(self.readAction)
        fileMenu.addAction(self.loadProjectAct)
        fileMenu.addAction(self.saveAction)

        runMenu = menubar.addMenu('&运行')
        runMenu.addAction(self.labelObject)
        runMenu.addAction(self.detectObject)
        runMenu.addAction(self.trainDetector)


    def keyPressEvent(self, ev):
        self.image1.keyPressEvent(ev)

    def runTrainer(self):
        if self.isDetecting == True :
            QMessageBox.information(self , "提示" ,
                                    "正在训练，请稍等")
            return

        if self.filedock1.filePath :
            # run_detector(self.filedock1.filePath)
            self.isDetecting = True

            self.prj_cfg['img_path'] = self.filedock1.filePath

            self.det = ChangeDetector(self.prj_cfg , self.label_status , "Train")
            self.det.trigger.connect(self.drawResult)
            self.det.start()


    def runDetector(self):
        if self.isDetecting == True:
            QMessageBox.information(self, "提示",
                                    "正在检测，请稍等")
            return


        if self.filedock1.filePath:
            #run_detector(self.filedock1.filePath)
            self.isDetecting = True

            self.prj_cfg['img_path'] = self.filedock1.filePath

            self.det = ChangeDetector(self.prj_cfg, self.label_status, "Inference")
            self.det.trigger.connect(self.drawResult)
            self.det.start()


    def drawResult(self):
        self.filedock1.filePath = os.path.join(sys_cfg['project_path'] , "resizedDetectedImg.jpg")
        self.newCanvas()

    def newProject(self):
        if self.enableCheckSave:
            self.checkSave()

        self.prj_cfg = Config()
        self.prj_cfg['labels'] = self.cfg['label_template']
        wizard = MagicWizard(self, self.prj_cfg)
        wizard.show()
        wizard.button(QWizard.FinishButton).clicked.connect(self.creatProject)


    def creatProject(self):
        print("self.prj_cfg:", self.prj_cfg)
        self.prj_cfg.save(os.path.join(self.prj_cfg['project_path'], self.prj_cfg['project_name']+'.yaml'))
        self.prj_cfg.cfg_merge(self.cfg)
        sys_cfg.cfg_merge(self.prj_cfg)

        self.label_color = []
        print("color", self.prj_cfg['label_color'])
        for  ik, iv in self.prj_cfg['label_color'].items():
            print(iv)
            color = QColor(iv[0], iv[1], iv[2])
            self.label_color.append(color)


        image_list = self.prj_cfg["image_list"]
        if len(image_list) > 0:
            # set image path & dir
            imgfp = image_list[0]
            self.filedock1.filePath = imgfp
            self.filedock1.directory = [os.path.dirname(imgfp)]

            # add images to file docker
            self.filedock1.clear()
            for imgf in image_list:
                self.filedock1.addFile(imgf)
            self.newCanvas()



    def loadProject(self):
        if self.enableCheckSave:
            self.checkSave()

        if self.isDetecting == True:
            QMessageBox.information(self, "提示",
                                    "正在检测，请稍后打开项目")
            return

        # select project file path
        prj_filePath = QFileDialog.getOpenFileNames(self,
                                                "选取项目文件",
                                                sys_cfg["system_rootpath"],
                                                "Project file (*.yaml)")[0][0]

        # load project config
        prj_cfg = Config()
        prj_cfg.parse_arguments(prj_filePath)
        prj_cfg["system_rootpath"] = sys_cfg["system_rootpath"]
        prj_cfg.cfg_merge(self.cfg)
        self.prj_cfg = prj_cfg
        sys_cfg.cfg_merge(prj_cfg)

        # set first image to self.filedocker1
        image_list = prj_cfg["image_list"]
        if len(image_list) > 0:
            # set image path & dir
            imgfp = image_list[0]
            self.filedock1.filePath = imgfp
            self.filedock1.directory = [os.path.dirname(imgfp)]

            # add images to file docker
            self.filedock1.clear()
            for imgf in image_list:
                self.filedock1.addFile(imgf)

            self.newCanvas()

    def labelImg(self):
        self.image1.toggleDrawMode(False , createMode="polygon")
        self.image1.update()

    def save(self):
        if self.isDetecting == True:
            QMessageBox.information(self, "提示",
                                    "正在检测，请稍后保存项目结果")
            return

        if self.filedock1.filePath is not None:
            self.enableCheckSave = False

            fp_base, fp_ext = os.path.splitext(self.filedock1.filePath)
            path = fp_base + ".lab"
            saveFile = open(path, 'w')

            img1 = cv2.imread(self.filedock1.filePath)
            h, w, _ = img1.shape

            for shape in self.v1:
                if shape.label in self.labelitem:
                    clc = self.labelitem.index(shape.label)

                    vert1, vert2 = shape[0], shape[2]
                    saveFile.write(str(clc) + ' ' + str(int(vert1.x())) + ' ' + str(int(vert1.y())) + ' ' + str(int(vert2.x())) + ' ' + str(int(vert2.y())) + '\n')
            saveFile.close()

            QMessageBox.information(self, "提示",
                "保存成功")

        else:
            QMessageBox.information(self, "提示",
                                    "请先新建文件")


    def checkSave(self):
        if self.v1:
            reply = QMessageBox.information(self,
                                            "提示",
                                            "当前界面未保存,是否保存?",
                                            QMessageBox.Yes | QMessageBox.No)
            if reply == QMessageBox.Yes:
                self.save()

    def fileDockerAddFile(self):
        filelist = QFileDialog.getOpenFileNames(self,
                                                "选取文件夹",
                                                sys_cfg["system_rootpath"],
                                                "".join(extensions))[0]
        print(filelist)

        for f in filelist:
            self.filedock1.addFile(f)

    def loadLabels(self, shapes):
        s = []
        for shape in shapes:
            label = "change"
            points = shape["points"]
            shape_type = "polygon"
            flags = "none"
            group_id = "none"
            other_data = "none"

            if not points:
                # skip point-empty shape
                continue

            shape = Shape(
                label=label,
                shape_type=shape_type,
                group_id=group_id,
            )
            for x, y in points:
                shape.addPoint(QPointF(x, y))
            shape.close()

            default_flags = {}
            shape.flags = default_flags
            shape.flags.update(flags)
            shape.other_data = other_data

            s.append(shape)
        self.loadShapes(s)

    def loadShapes(self, shapes, replace=True):
        self.image1.loadShapes(shapes, replace=replace)


    def newCanvas(self):
        # self.createcenter(self.grid)
        if self.enableCheckSave:
            self.checkSave()
        self.enableCheckSave = False
        self.v1.clear()

        self.cleanTmp()

        self.scroll1.zoom = 100
        self.image1.loadPixmap(QPixmap(self.filedock1.filePath))
        sys_cfg["current_img_path"] = self.filedock1.filePath



        # load change detected segmentation
        cur_img_path = sys_cfg["current_img_path"]
        basename = os.path.basename(cur_img_path)
        basename = basename.rsplit('.' , 1)[0]
        labfp = os.path.join(sys_cfg['project_path'], basename+".json")



        if os.path.exists(labfp):
            with open(labfp , "r") as f :
                data = json.load(f)

            # self.image1.shapes = data["shapes"]
            shapes = [
                dict(
                    label="change" ,
                    points=s["points"] ,
                    shape_type=s.get("shape_type" , "polygon") ,
                    flags=s.get("flags" , {}) ,
                    group_id=s.get("group_id") ,
                    other_data="none",
                )
                for s in data["shapes"]
            ]

            s = []
            for shape in shapes :
                label = shape["label"]
                points = shape["points"]
                shape_type = shape["shape_type"]
                flags = shape["flags"]
                group_id = shape["group_id"]
                other_data = shape["other_data"]

                if not points :
                    # skip point-empty shape
                    continue

                shape = Shape(
                    label=label ,
                    shape_type=shape_type ,
                    group_id=group_id ,
                )
                for x , y in points :
                    shape.addPoint(QPointF(x , y))
                shape.close()

                default_flags = {}
                shape.flags = default_flags
                shape.flags.update(flags)
                shape.other_data = other_data

                s.append(shape)

            self.loadShapes(s)
            # self.image1.shapes = shapes


        # labfp = os.path.splitext(self.filedock1.filePath)[0] + ".lab"
        # if os.path.exists(labfp) :
        #     self.v1.clear()
        #
        #     fp = open(labfp)
        #     self.addLabelList()  # FIXME: need set labels from prj_cfg
        #     for lines in fp.readlines() :
        #         clc , left , top , right , bot = lines.strip().split()
        #         clc , left , right , top , bot = int(clc) , int(left) , int(right) , int(top) , int(bot)
        #         current = Shape(self.labelitem[clc] , self.label_color[clc])
        #         current.addPoint(QPointF(left , top)) , current.addPoint(QPointF(right , top))
        #         current.addPoint(QPointF(right , bot)) , current.addPoint(QPointF(left , bot))
        #         self.v1.append(current)
        #     fp.close()
        #     self.isDetecting = False
        if self.filedock1.filePath:
            from PIL import Image
            im = Image.open(self.filedock1.filePath)  #fixme
            img_size = max(im.size[0],im.size[1])
            self.image1.base_epsilon = img_size / 20
            self.image1.epsilon = self.image1.base_epsilon
            self.scroll1.scaleFitWindow()

        self.isDetecting = False
        self.update()

    def closeEvent(self, Event):
        self.cleanTmp()



    def cleanTmp(self):
        def deletetmp():
            if os.path.exists('tmp/'):
                shutil.rmtree('tmp/', True)

        threads = []
        for i in range(10):
            t = threading.Thread(target=deletetmp)
            threads.append(t)
            t.start()
        for t in threads:
            t.join()


list_food = ["Pizza", "Taco", "Burrito"]

class Demo(QWidget):
    def __init__(self, parent):
        super(Demo, self).__init__()
        #self.resize(80,10)


    def setCenter(self):
        screen = QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move(int((screen.width() - size.width()) / 2) , int((screen.height() - size.height()) / 2))

    def initUI(self, list_label):
        self.groupBox = QGroupBox("请选择标注的标签")

        # Checkbox layout
        self.hbox = QHBoxLayout()

        self.chk_lst = []
        for i in range(len(list_label)):
            self.cbx = QRadioButton(list_label[i])
            self.chk_lst.append(self.cbx)
            self.hbox.addWidget(self.cbx)
            self.cbx.toggled.connect(lambda :self.btnstate(self.sender().text(), self.sender().isChecked()))
        self.chk_lst[0].setChecked(True)
        self.groupBox.setLayout(self.hbox)

        # vertical box layout
        self.vlayout = QVBoxLayout()
        self.vlayout.addWidget(self.groupBox)
        self.vlayout.addStretch()
        self.setLayout(self.vlayout)


    def btnstate(self, btn, flag):
        if flag:
            self.currentLabel = btn